#include <iostream>
#include <opencv2/opencv.hpp>
#include <opencv2/dnn.hpp>
#include <windows.h> // 包含 Windows API 的头文件
#include <onnxruntime_cxx_api.h>

using namespace std;
using namespace cv;


int main() {
    SetConsoleCP(65001); // 设置控制台输入为 UTF-8 编码
    SetConsoleOutputCP(65001); // 设置控制台输出为 UTF-8 编码
    printf("hello world!\n");

    // 创建 ONNX Runtime 环境
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "ONNXRuntimeTest");

    // 创建 SessionOptions
    Ort::SessionOptions session_options;
    session_options.SetIntraOpNumThreads(1);  // 设置线程数
    session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);

    // 模型路径
    const wchar_t *model_path = L"D:\\workspaces\\opencv_test\\onnx_test\\cmake-build-debug\\efficient_sam_vitt.onnx";  // 替换为你的 ONNX 模型路径

    // 创建 Session
    Ort::Session session(env, reinterpret_cast<const wchar_t *>(model_path), session_options);

    // 输出模型信息
    std::cout << "Model loaded successfully!" << std::endl;


    // 加载图片
    std::string imagePath = "D:\\workspaces\\opencv_test\\onnx_test\\dogs.jpg"; // 替换为你的图片路径
    // 判断路径是否存在
    if (imagePath.empty()) {
        std::cout << "图片路径不存在" << std::endl;
        return -1;
    }
    else {
        std::cout << "图片路径存在" << std::endl;
    }

    cv::Mat image = cv::imread(imagePath);
    if (image.empty()) {
        std::cerr << "Failed to load image!" << std::endl;
        return -1;
    }

    // 输入点，形状为 batch_size x num_queries x num_points x 2
    std::vector<std::vector<std::vector<std::array<float, 2>>>> input_points = {
            {{{580.0f, 350.0f}, {650.0f, 350.0f}}}
    };

    // 输入标签，形状为 batch_size x num_queries x num_points
    std::vector<std::vector<std::vector<float>>> input_labels = {
            {{1.0f, 1.0f}}
    };

    // 获取输入/输出节点信息
    Ort::AllocatorWithDefaultOptions allocator;
    std::string input_name0 = session.GetInputNameAllocated(0, allocator).get();  // "batched_images"
    std::string input_name1 = session.GetInputNameAllocated(1, allocator).get();  // "batched_point_coords"
    std::string input_name2 = session.GetInputNameAllocated(2, allocator).get();  // "batched_point_labels"
    std::cout << "Input names: " << input_name0 << ", " << input_name1 << ", " << input_name2 << std::endl;

    // 获取输出节点
    std::string output_name = session.GetOutputNameAllocated(0, allocator).get();
    std::vector<const char*> output_names = {output_name.c_str()};
    std::cout << "Output name: " << output_name << std::endl;

    // 预处理输入图像: batch, channels, height, width
    cv::Mat input_image = image.clone();
    input_image.convertTo(input_image, CV_32F, 1.0 / 255.0);
    cv::Mat input_data = cv::dnn::blobFromImage(
            input_image,
            1.0,
            cv::Size(),
            cv::Scalar(0, 0, 0),
            true,
            false
        );
    std::cout << "input_data shape: " << input_data.size << std::endl;

//    return 0;

    // 预处理输入点和标签
    std::vector<float> input_points_data;
    std::vector<float> input_labels_data;
    for (const auto &batch : input_points) {
        for (const auto &query : batch) {
            for (const auto &point : query) {
                input_points_data.push_back(point[0]);
                input_points_data.push_back(point[1]);
            }
        }
    };

    for (const auto &batch : input_labels) {
        for (const auto &query : batch) {
            for (const auto &label : query) {
                input_labels_data.push_back(label);
            }
        }
    };


    // 准备输入数据
    std::vector<int64_t> input_shape0 = {1, 3, image.rows, image.cols};
    std::vector<int64_t> input_shape1 = {1, 1, 2, 2};
    std::vector<int64_t> input_shape2 = {1, 1, 2};

    // 创建 ONNX 输入张量
    Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);

    // 输入张量
    Ort::Value input_tensor0 = Ort::Value::CreateTensor<float>(memory_info, input_data.ptr<float>(), input_data.total(), input_shape0.data(), input_shape0.size());
    Ort::Value input_tensor1 = Ort::Value::CreateTensor<float>(memory_info, input_points_data.data(), input_points_data.size(), input_shape1.data(), input_shape1.size());
    Ort::Value input_tensor2 = Ort::Value::CreateTensor<float>(memory_info, input_labels_data.data(), input_labels_data.size(), input_shape2.data(), input_shape2.size());
    // 整理输入
    // 使用 emplace_back 显式添加到 vector
    std::vector<Ort::Value> input_tensors;
    input_tensors.emplace_back(std::move(input_tensor0));
    input_tensors.emplace_back(std::move(input_tensor1));
    input_tensors.emplace_back(std::move(input_tensor2));

    Ort::RunOptions run_options; // 创建运行选项
    std::vector<const char*> input_names = {input_name0.c_str(), input_name1.c_str(), input_name2.c_str()};
    // 运行推理
    auto start_time = std::chrono::high_resolution_clock::now();
    auto output_tensors = session.Run(run_options, input_names.data(), input_tensors.data(), input_tensors.size(), output_names.data(), output_names.size());
    auto end_time = std::chrono::high_resolution_clock::now();
    // 计算推理时间
    auto inference_time = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time);
    std::cout << "Inference time: " << inference_time.count() << " ms" << std::endl;

    // 获取输出结果
    auto output_data = output_tensors[0].GetTensorMutableData<float>();
    // 输出结果形状
    auto output_shape = output_tensors[0].GetTensorTypeAndShapeInfo().GetShape();
    // 输出结果
    std::cout << "Output shape: ";
    for (auto dim : output_shape) {
        std::cout << dim << " ";
    }
    std::cout << std::endl;

    // 显示结果
    int output_height = static_cast<int>(output_shape[4]);
    int output_width = static_cast<int>(output_shape[3]);
    std::vector<float> masks(output_data, output_data + output_height * output_width);

    cv::Mat mask_image(output_width, output_height, CV_32F, (void*)masks.data());
    std::cout << "mask_image shape: " << mask_image.size << std::endl;
    // 可视化结果
    cv::Mat visualized_mask;
    cv::threshold(mask_image, visualized_mask, 0, 255, cv::THRESH_BINARY);
    visualized_mask.convertTo(visualized_mask, CV_8U); // 转换为 8 位
    cv::cvtColor(visualized_mask, visualized_mask, cv::COLOR_GRAY2BGR);

    cv::Mat result;
    cv::hconcat(image, visualized_mask, result); // 水平拼接 (水平方向)
    
    cv::imshow("Result", result);
    cv::waitKey(0);
    
    return 0;
}
